# -*- coding: utf-8 -*-
"""
@Time    : 2025/1/23 11:41 
@Author  : ZhangShenao 
@File    : agent.py 
@Desc    : Agent推理框架
"""
import os

import dotenv
from openai import OpenAI

DEFAULT_MODEL = "gpt-4o-mini"

# 初始化OpenAI客户端
dotenv.load_dotenv()
client = OpenAI(base_url=os.getenv("OPENAI_API_BASE"))


class Agent:
    """Agent推理框架"""

    def __init__(self, system_prompt=""):
        """初始化方法"""

        # 初始化消息列表
        self._messages = []
        # 构造SystemPrompt
        if system_prompt is not None and system_prompt != "":
            self._messages.append({"role": "system", "content": system_prompt})

    def invoke(self, query: str) -> str:
        """调用Agent"""

        # 在Prompt中添加消息列表
        self._messages.append({"role": "user", "content": query})
        # 执行推理,获取结果
        result = self.exec()
        # 保存聊天历史
        self._messages.append({"role": "assistant", "content": result})

        # 返回结果
        return result

    def exec(self) -> str:
        """执行推理,返回结果"""

        # 调用OpenAI接口
        completion = client.chat.completions.create(
            model=DEFAULT_MODEL,
            messages=self._messages,
            temperature=0,  # 设置较小的temperature,因为是推理的场景,需要生成更确定性的结果
        )

        # 返回结果
        return completion.choices[0].message.content
